"""
    @file Sklearn
    @description Decision Tree By Sklearn Using Gini Index
    @author Synhard
    @tel 13001321080
    @id 21126338
    @email 823436512@qq.com
    @date 2021-09-25 14:43
    @version 1.0
"""
import xlrd
import numpy as np
from collections import Counter
import matplotlib.pyplot as plt
import matplotlib
from sklearn.metrics import confusion_matrix

# 可视化部分
matplotlib.rcParams['font.family'] = 'SimHei'  # 用来正常显示中文
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

decisionNode = dict(boxstyle="sawtooth", fc="0.8")
leafNode = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")


def getNumLeafs(myTree):
    numLeafs = 0
    firstStr = list(myTree.keys())[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[
                    key]).__name__ == 'dict':
            numLeafs += getNumLeafs(secondDict[key])
        else:
            numLeafs += 1
    return numLeafs


def getTreeDepth(myTree):
    maxDepth = 0
    firstStr = list(myTree.keys())[0]
    secondDict = myTree[firstStr]
    for key in secondDict.keys():
        if type(secondDict[
                    key]).__name__ == 'dict':
            thisDepth = 1 + getTreeDepth(secondDict[key])
        else:
            thisDepth = 1
        if thisDepth > maxDepth: maxDepth = thisDepth
    return maxDepth


def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    createPlot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',
                            xytext=centerPt, textcoords='axes fraction',
                            va="center", ha="center", bbox=nodeType, arrowprops=arrow_args)


def plotMidText(cntrPt, parentPt, txtString):
    xMid = (parentPt[0] - cntrPt[0]) / 2.0 + cntrPt[0]
    yMid = (parentPt[1] - cntrPt[1]) / 2.0 + cntrPt[1]
    createPlot.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)


def plotTree(myTree, parentPt, nodeTxt):
    numLeafs = getNumLeafs(myTree)
    firstStr = list(myTree.keys())[0]
    cntrPt = (plotTree.xOff + (2 + float(numLeafs)) / 2.0 / plotTree.totalW, plotTree.yOff)
    plotMidText(cntrPt, parentPt, nodeTxt)
    plotNode(firstStr, cntrPt, parentPt, decisionNode)
    secondDict = myTree[firstStr]
    plotTree.yOff = plotTree.yOff - 1.0 / plotTree.totalD
    for key in secondDict.keys():
        if type(secondDict[
                    key]).__name__ == 'dict':
            plotTree(secondDict[key], cntrPt, str(key))
        else:
            plotTree.xOff = plotTree.xOff + 1.1 / plotTree.totalW
            plotNode(secondDict[key], (plotTree.xOff, plotTree.yOff), cntrPt, leafNode)
            plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(key))
    plotTree.yOff = plotTree.yOff + 1.0 / plotTree.totalD


def createPlot(myTree):
    fig = plt.figure(1, facecolor='white', figsize=(8, 8))
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    createPlot.ax1 = plt.subplot(111, frameon=False, **axprops)
    plotTree.totalW = float(getNumLeafs(myTree))
    plotTree.totalD = float(getTreeDepth(myTree))
    plotTree.xOff = -1.5 / plotTree.totalW
    plotTree.yOff = 1.1
    plotTree(myTree, (0.5, 1.0), '')
    plt.show()


def read_xslx(xslx_path):
    trainingdata = []  # 先声明一个空list
    testingdata = []
    data = xlrd.open_workbook(xslx_path)  # 读取文件
    table = data.sheet_by_index(0)  # 按索引获取工作表，0就是工作表1

    for i in range(int(0.7 * table.nrows)):  # table.nrows表示总行数
        line = table.row_values(i)  # 读取每行数据，保存在line里面，line是list
        trainingdata.append(line)  # 将line加入到trainingdata中，trainingdata是二维list
    trainingdata = np.array(trainingdata)  # 将trainingdata从二维list变成数组

    testingdata.append(table.row_values(0))
    for i in range(int(0.7 * table.nrows), int(table.nrows)):  # table.nrows表示总行数
        line = table.row_values(i)  # 读取每行数据，保存在line里面，line是list
        testingdata.append(line)  # 将line加入到testingdata中，testingdata是二维list
    testingdata = np.array(testingdata)  # 将testingdata从二维list变成数组

    return trainingdata, testingdata


# 计算结点GINI值
def calcGini(dataSet):
    numTotal = dataSet.shape[0]  # 记录本数据集总条数
    length = len(dataSet[0])  # 计算特征列数
    frequent_0 = 0.0  # 记录两种样本出现次数
    frequent_1 = 0.0
    for i in range(0, numTotal):
        if dataSet[i][length - 1] == '0.0':
            frequent_0 += 1
        elif dataSet[i][length - 1] == '1.0':
            frequent_1 += 1
    gini = 1 - (frequent_0 / numTotal) ** 2 - (frequent_1 / numTotal) ** 2
    return gini


# 根据条件分离数据集
def splitDataSet(dataSet, n, value, type):
    subDataSet = []
    numTotal = dataSet.shape[0]  # 记录本数据集总条数
    if type == 1:  # type==1对应小于等于value的情况
        for i in range(0, numTotal):
            if float(dataSet[i][n]) <= float(value):
                subDataSet.append(dataSet[i])
    elif type == 2:  # type==2对应大于value的情况
        for i in range(0, numTotal):
            if float(dataSet[i][n]) > float(value):
                subDataSet.append(dataSet[i])
    subDataSet = np.array(subDataSet)  # 强制转换为array类型
    pass
    return subDataSet, len(subDataSet)


# 选择最好的特征划分数据集，即返回最佳特征下标及传入数据集各列的Gini指数
def FindBestFeature(dataSet):
    numTotal = dataSet.shape[0]  # 记录本数据集总条数
    numFeatures = len(dataSet[0]) - 1  # 计算特征列数
    bestFeature = -1  # 初始化参数，记录最优特征列i，下标从0开始
    columnFeaGini = {}  # 初始化参数，记录每一列x的每一种特征的基尼 Gini(D,A)
    for i in range(0, numFeatures - 1):  # 遍历所有x特征列,i为特征标号
        featList = list(dataSet[:, i])  # 取这一列x中所有数据，转换为list类型
        featListSort = [float(x) for x in featList]
        featListSort.sort()  # 对该特征值排序
        FeaGinis = []
        FeaGiniv = []
        for j in range(0, len(featListSort) - 1):  # j为第几组数据
            value = (featListSort[j] + featListSort[j + 1]) / 2
            feaGini = 0.0
            subDataSet1, sublen1 = splitDataSet(dataSet, i, value, 1)  # 获取切分后的数据
            subDataSet2, sublen2 = splitDataSet(dataSet, i, value, 2)
            if sublen1 == 0:
                feaGini = (sublen2 / numTotal) * calcGini(subDataSet2)  # 计算此分法对应Gini值
            elif sublen2 == 0:
                feaGini = (sublen1 / numTotal) * calcGini(subDataSet1)
            else:
                feaGini = (sublen1 / numTotal) * calcGini(subDataSet1) + (sublen2 / numTotal) * calcGini(
                    subDataSet2)  # 计算此分法对应Gini值
            FeaGinis.append(feaGini)  # 记录该特征下各种分法遍历出的Gini值
            FeaGiniv.append(value)  # 记录该特征下的各种分法

        columnFeaGini['%d_%f' % (i, FeaGiniv[FeaGinis.index(min(FeaGinis))])] = min(FeaGinis)  # 将该特征下最小的Gini值
    bestFeature = min(columnFeaGini, key=columnFeaGini.get)  # 找到最小的Gini指数对应的数据列
    return bestFeature, columnFeaGini


# 生成决策树。输入：训练数据集D，特征集A。输出：决策树T
def createTree(dataSet, features, decisionTree):
    if len(features) > 2:  # 特征未用完
        bestFeature, columnFeaGini = FindBestFeature(dataSet)
        bestFeatureLable = features[int(bestFeature.split('_')[0])]  # 最佳特征
        NodeName = bestFeatureLable + '\n' + '<=' + bestFeature.split('_')[1]  # 结点名称
        decisionTree = {NodeName: {}}  # 构建树，以Gini指数最小的特征bestFeature为子节点
    else:
        return decisionTree

    LeftSet, LeftSet_len = splitDataSet(dataSet, int(bestFeature.split('_')[0]), float(bestFeature.split('_')[1]), 1)
    RightSet, RightSet_len = splitDataSet(dataSet, int(bestFeature.split('_')[0]), float(bestFeature.split('_')[1]), 2)
    del (features[int(bestFeature.split('_')[0])])  # 该特征已为子节点使用，则删除，以便接下来继续构建子树

    if calcGini(LeftSet) <= 0.1 or len(features) == 2:
        L_lables_grp = dict(Counter(LeftSet[:, -1]))
        L_leaf = max(L_lables_grp, key=L_lables_grp.get)  # 获得划分后出现概率最大的分类作为结点的分类
        decisionTree[NodeName]['Y'] = L_leaf  # 设定左枝叶子值
    elif calcGini(LeftSet) > 0.1:
        dataSetNew = np.delete(LeftSet, int(bestFeature.split('_')[0]), axis=1)  # 删除此最优划分x列，使用剩余的x列进行数据划分
        L_subFeatures = features[:]
        decisionTree[NodeName]['Y'] = {'NONE'}
        decisionTree[NodeName]['Y'] = createTree(dataSetNew, L_subFeatures, decisionTree[NodeName]['Y'])  # 递归生成左边的树

    if calcGini(RightSet) <= 0.1 or len(features) == 2:
        R_lables_grp = dict(Counter(RightSet[:, -1]))
        R_leaf = max(R_lables_grp, key=R_lables_grp.get)  # 获得划分后出现概率最大的分类作为结点的分类
        decisionTree[NodeName]['N'] = R_leaf  # 设定右枝叶子值
    elif calcGini(RightSet) > 0.1:
        dataSetNew = np.delete(RightSet, int(bestFeature.split('_')[0]), axis=1)  # 删除此最优划分x列，使用剩余的x列进行数据划分
        R_subFeatures = features[:]
        decisionTree[NodeName]['N'] = {'NONE'}
        decisionTree[NodeName]['N'] = createTree(dataSetNew, R_subFeatures, decisionTree[NodeName]['N'])  # 递归生成右边的树

    return decisionTree


# 获得测试结果
def testTree(decisontree, dataSet):
    testmemory = []
    label = []
    TP = 0
    FP = 0
    TN = 0
    FN = 0
    rowDict = []
    for i in range(1, dataSet.shape[0]):
        dict = {}
        for j in range(0, dataSet[0].shape[0]):
            dict[dataSet[0][j]] = dataSet[i][j]

        rowDict.append(dict)

    for row in rowDict:
        label.append(row['target'])
        testmemory.append(prediction(decisontree, row))
        pass

    for i in range(0, len(rowDict)):
        if (testmemory[i] == '1.0') and (label[i] == '1.0'):  # test为预测 label为实际
            TP += 1
        elif (testmemory[i] == '1.0') and (label[i] == '0.0'):
            FP += 1
        elif (testmemory[i] == '0.0') and (label[i] == '0.0'):
            TN += 1
        elif (testmemory[i] == '0.0') and (label[i] == '1.0'):
            FN += 1

    print('TP:%d' % TP)  # 真阳性
    print('FP:%d' % FP)  # 假阳性
    print('TN:%d' % TN)  # 真阴性
    print('FN:%d' % FN)  # 假阴性

    cm = confusion_matrix(label, testmemory, labels=["0.0", "1.0"])
    plt.rc('figure', figsize=(8, 8))
    plt.matshow(cm, cmap=plt.cm.cool)  # 背景颜色
    plt.colorbar()  # 颜色标签
    # 内部添加图例标签
    for x in range(len(cm)):
        for y in range(len(cm)):
            plt.annotate(cm[x, y], xy=(y, x), horizontalalignment='center', verticalalignment='center')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.title('decision_tree')
    plt.savefig(r'confusion_matrix')
    plt.show()


def prediction(decisontree, row):
    if isinstance(decisontree, str):
        return decisontree
    root = list(decisontree.keys())[0]
    attr = root.splitlines()[0]
    value = float(str(root.splitlines()[1])[2:])
    if float(row.get(attr)) <= value:
        return prediction(decisontree.get(root).get('Y'), row)
    else:
        return prediction(decisontree.get(root).get('N'), row)
    pass


if __name__ == '__main__':
    trainingData, testingData = read_xslx(r'heart.xlsx')
    features = list(trainingData[0])  # 表头
    trainingDataSet = trainingData[1:]  # 训练集
    bestFeature, columnFeaGini = FindBestFeature(trainingDataSet)
    decisionTree = createTree(trainingDataSet, features, {})  # 建立决策树，CART分类树
    testTree(decisionTree, testingData)
    createPlot(decisionTree)
